import tensorflow as tf
from getStrokeLabel import genStrokes
#from getStroke import genStrokes
import numpy as np
from config import Config
from myCNN import CNN2
import os
import h5py
import matplotlib.pyplot as plt
import cv2


def read_data(path):
     images = np.array(h5py.File(path, 'r').get('images'))
     labels = np.array(h5py.File(path, 'r').get('labels'))
     return images, labels

def read_matte(path):
     mattes = np.array(h5py.File(path, 'r').get('matte'))
     return mattes


config = Config()
n_steps = config.step
draw = genStrokes(config, )
cnn2 = CNN2(config,)


def MSE(image,gt):
      if len(image.shape)>2:
           image=image[:,:,0]
      if len(gt.shape)>2:
           gt=gt[:,:,0]
  
      width=image.shape[0]
      height=image.shape[1]
      sumup=0
      for j in range(width):
         for k in range(height):
             sumup=sumup+(image[j,k]-gt[j,k])*(image[j,k]-gt[j,k])
      MSE=sumup/(width*height)
      return MSE
    
matte_dir = os.path.join(os.getcwd(), 'matte.h5')
data_dir = os.path.join(os.getcwd(), 'train.h5')
train_images, train_alphas = read_data(data_dir)
matte_arr=read_matte(matte_dir)
strokeMaps = train_images.copy()
images=train_images.copy()



this_matte=matte_arr[-1,0,:,:]
gt=train_alphas[0]
image=np.float32(cv2.imread('./first_try_DATA/tmp/matte.png'))
print MSE(image,gt)
cv2.imshow
print MSE(this_matte,gt)
